package com.deepseekrag.controller;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;


@RestController
@RequestMapping("/api/v1/ollama")
public class OllamaController {
    @Resource
    private ChatClient chatClient;

    @Resource
    private ObjectMapper objectMapper;

    @GetMapping("/generate")
    public String generate(@RequestParam("message") String message) {
        return chatClient.prompt()
                .user(message)
                .call()
                .content();
    }

    @GetMapping(value = "/generate_stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> generateStream(@RequestParam("message") String message) {
        return chatClient.prompt()
                .user(message)
                .stream()
                .chatResponse()
                .map(chatResponse -> {
                    try {
                        String json = objectMapper.writeValueAsString(chatResponse);
                        return ServerSentEvent.<String>builder().data(json).build();
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException("Failed to serialize ChatResponse", e);
                    }
                });
    }

}
